import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils import data
import torchvision.transforms as transforms
from PIL import Image
import random
import numpy as np
import cv2
from PIL import Image
import math

class TrainDataset(data.Dataset):
    def __init__(self, data_root, data_list, weight, height):
        self.weight = weight
        self.height = height
        self.sal_root = data_root
        self.sal_source = data_list

        with open(self.sal_source, 'r') as f:
            self.sal_list = [x.strip() for x in f.readlines()]

        self.sal_num = len(self.sal_list)

        self.img_transform = transforms.Compose([
            transforms.Resize((self.weight, self.height)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop((self.weight, self.height), padding=4),
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])])

        self.gt_transform = transforms.Compose([
            transforms.Resize((self.weight, self.height)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop((self.weight, self.height), padding=4),
            transforms.ToTensor()])


    def __getitem__(self, index):
        im_name = self.sal_list[index % self.sal_num].split()[0]
        gt_name = self.sal_list[index % self.sal_num].split()[1]
        # print(self.sal_root+im_name, self.sal_root+gt_name)
        sal_image = self.load_img(self.sal_root + im_name)
        sal_label = self.load_gt(self.sal_root + gt_name)

        seed = torch.randint(0, 2**32, (1,)).item()
        torch.manual_seed(seed)
        sal_image = self.img_transform(sal_image)

        torch.manual_seed(seed)
        sal_label = self.gt_transform(sal_label)

        return sal_image, sal_label

    def __len__(self):
        return self.sal_num

    def load_img(self, path):
        with open(path, 'rb') as f:
            img = Image.open(f)
            return img.convert('RGB')

    def load_gt(self, path):
        with open(path, 'rb') as f:
            img = Image.open(f)
            return img.convert('L')


class TestDataset(data.Dataset):
    def __init__(self, data_root, data_list, weight, height):
        self.weight = weight
        self.height = height

        self.sal_root = data_root
        self.sal_source = data_list
        with open(self.sal_source, 'r') as f:
            self.sal_list = [x.strip() for x in f.readlines()]
        self.sal_num = len(self.sal_list)

        self.img_transform = transforms.Compose([
            transforms.Resize((self.weight, self.height)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])])

        self.gt_transform = transforms.Compose([transforms.ToTensor()])

    def __getitem__(self, index):
        im_name = self.sal_list[index % self.sal_num].split()[0]
        gt_name = self.sal_list[index % self.sal_num].split()[1]
        # print(self.sal_root+im_name, self.sal_root+gt_name)
        sal_image = self.load_img(self.sal_root + im_name)
        sal_label = self.load_gt(self.sal_root + gt_name)
        sal_image = self.img_transform(sal_image)
        sal_label = self.gt_transform(sal_label)
        return sal_image,sal_label,gt_name

    def __len__(self):
        return self.sal_num

    def load_img(self, path):
        with open(path, 'rb') as f:
            img = Image.open(f)
            return img.convert('RGB')

    def load_gt(self, path):
        with open(path, 'rb') as f:
            img = Image.open(f)
            return img.convert('L')


def get_train_loader(data_root, data_list, weight=352, height=352, batch_size=1):
    set1 = TrainDataset(data_root, data_list, weight, height)
    loader = DataLoader(set1, batch_size=batch_size, shuffle=True)
    return loader

def get_test_loader(data_root, data_list, weight=352, height=352, batch_size=1):
    set1 = TestDataset(data_root, data_list, weight, height)
    loader = DataLoader(set1, batch_size=batch_size, shuffle=True)
    return loader
